
anqar_panel = readRDS('data/anqar_panel_2023.rds')
df <- readRDS('data/full_panel_df.rds') %>% 
  filter(year >= 2008) %>% 
  filter(year <= 2014) %>% 
  mutate(
    cohort = ifelse(cohort == '10000', '0', cohort)
  ) %>% 
  select('DISTID', 'year', 'cohort',  'treated_ever', 
         'aog_no_ied', 'ied', 'IED_Explosion', 'DF', 'totalcas'
  ) 
df = merge(df, anqar_panel, by = c('DISTID', 'year'), all.x = T)

force_balance = function(data, outcome, baseline_year){
  
  data = data 
  data$outcome = outcome 
  
  data <- data %>% 
    group_by(DISTID) %>% 
    filter(year >= baseline_year) %>% 
    mutate(na_ever = mean(outcome)) %>% 
    filter(is.na(na_ever) == F) %>% 
    select(-c(na_ever)) %>% 
    as.data.frame()
  
  return(data)
}

tjbal_unbalance = function(data, outcome, baseline_year, 
                           X,
                           estimator_type, vce){
  
  force_balance = function(data, outcome, baseline_year){
    
    data = data 
    data$outcome = outcome 
    
    data <- data %>% 
      group_by(DISTID) %>% 
      filter(year >= baseline_year) %>% 
      mutate(na_ever = mean(outcome)) %>% 
      filter(is.na(na_ever) == F) %>% 
      select(-c(na_ever)) %>% 
      as.data.frame()
    
    return(data)
  }
  
  data = force_balance(data, outcome, baseline_year)
  
  m = tjbal(data, Y = 'outcome', 
            D = "treat", 
            index = c("DISTID", "year"), 
            X = X,
            demean = F, estimator = estimator_type, 
            vce = vce)
  
  return(m)
}


df = df %>% 
  mutate(
    cohort = ifelse(cohort == 'wd13', 2011, cohort), 
    
    cohort = ifelse(cohort == 'wd12', 2011, cohort), 
    
    cohort = ifelse(cohort == 'wd12est13', 2011, cohort), 
    
    cohort = ifelse(cohort == '0', 0, cohort), 
    
    treated = ifelse(cohort >0, as.numeric(year >= cohort, 1,0), 0), 
    
    treat = ifelse(cohort >0, as.numeric(year >= cohort, 1,0), 0)
    
  )


#### opinion #### 

out1 = tjbal_unbalance(data = df, outcome = df$dispute_court, 
                       baseline_year = 2009, 
                       X = NULL, 
                       estimator_type = "mean", 
                       vce = "jackknife")

out2 = tjbal_unbalance(data = df, outcome = df$return_good, 
                       baseline_year = 2009, 
                       X = NULL, 
                       estimator_type = "mean", 
                       vce = "jackknife")

out3 = tjbal_unbalance(data = df, outcome = df$gov_influence, 
                       baseline_year = 2008, 
                       X = NULL, 
                       estimator_type = "mean", 
                       vce = "jackknife")

out4 = tjbal_unbalance(data = df, outcome = df$gov_index2, 
                       baseline_year = 2008, 
                       X = NULL, 
                       estimator_type = "mean", 
                       vce = "jackknife")

att_tjbal_opinion <- matrix(, nrow = 4, ncol = 3)
out_list = rbind.data.frame(mod_1 = print(out1), mod_2 = print(out2), 
                            mod_3 = print(out3), mod_4 = print(out4))
out_list = out_list[ ,c(1,2,6)]

make_tjbal_table(out_list, 2, 1, out_names = c('State Court', 'Taliban Approval', 
                                               "Gov. Influence", 
                                               "Gov. Index"), 
                 data = df, id = df$DISTID, 
                 time = df$year, treated_group = df$treated_ever, 
                 outcome_matrix = df[,10:13], header = FALSE, header_labels = NA, 
                 label_treatment = 'Taliban Courts', extra_arg = NULL,
                 file_save = paste0(getwd(), "//tables", "//opinion_table_WD_recode.tex")
)


force_balance(df, df$dispute_court, 2009) %>% 
  group_by(year) %>% 
  summarise(n = n()) %>% 
  filter(year == min(year))
force_balance(df, df$gov_index2, 2008) %>% 
  group_by(year) %>% 
  summarise(n = n()) %>% 
  filter(year == min(year))


#### Combat #### 

df3 = df %>% 
  filter(year < 2014) %>% 
  dplyr::select('aog_no_ied', 'ied', "DF", 
                "IED_Explosion", 'totalcas',
                'treat', 'DISTID', 'year', 
                'cohort', 'treated_ever') %>% 
  as.data.frame()

J <- 5
out_names <- names(df3)[1:J]
att_tjbal <- matrix(, nrow = J, ncol = 3)

for(j in 1:J){
  
  att_tjbal[j, 1:3] <- print(tjbal(data = df3, Y = out_names[j], D = "treat", 
                                   index = c("DISTID", "year"), 
                                   demean = F, estimator = "mean", 
                                   vce = "jackknife"))[c(1,2,6)]
  
}

make_tjbal_table(att_tjbal, 2, 1, out_names = c('AOG', 'IED', "DF", 
                                                "IED Explosions", 'Causality Events'), 
                 data = df3, id = df3$DISTID, 
                 time = df3$year, treated_group = df3$treated_ever, 
                 outcome_matrix = df3[,1:5], header = FALSE, header_labels = NA, 
                 label_treatment = 'Taliban Courts', extra_arg = c('Dataset', 
                                                                   rep('ANSO', 2), 
                                                                   rep('SIGACT', 3)),
                 file_save = paste0(getwd(), "//tables", "//combat_table_WD_recode.tex")
)